class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def countUnivalSubtrees(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        self.value = 0
        if not root:
            return 0

        def search(root):
            if root.left:
                if root.right:
                    left_flag, left_value = search(root.left)
                    right_flag, right_value = search(root.right)
                    if left_flag and right_flag:
                        if left_value == right_value == root.val:
                            self.value += 1
                            return True, left_value
                        else:
                            return False, 0
                    else:
                        return False, 0
                else:
                    left_flag, left_value = search(root.left)
                    if left_flag:
                        if left_value == root.val:
                            self.value += 1
                            return True, left_value
                        else:
                            return False, 0
                    else:
                        return False, 0
            else:
                if root.right:
                    right_flag, right_value = search(root.right)
                    if right_flag:
                        if right_value == root.val:
                            self.value += 1
                            return True, right_value
                        else:
                            return False, 0
                    else:
                        return False, 0
                else:
                    self.value += 1
                    return True, root.val

        search(root)
        return self.value
